import random
from tqdm import tqdm
import numpy as np
from copy import deepcopy

def compute_hamming_dist_ordinal(x1, x2, normalize=False):
    hamming = (x1 != x2).sum()
    if normalize:
        return hamming / len(x1)
    return hamming

def sample_neighbour_ordinal(x, alphabet):
    x_pert = deepcopy(x)
    # Chooose a variable to modify
    choice = random.randint(0, len(x) - 1)
    # Obtain the current value.
    curr_val = x[choice]
    options = [i for i in range(alphabet) if i != curr_val]
    x_pert[choice] = random.choice(options)
    return x_pert

def local_search(x_center, f: callable,
                 max_hamming_dist,
                 alphabet,
                 n_restart: int = 1,
                 step: int = 100):
    """
    Local search algorithm
    :param n_restart: number of restarts
    :param x_center: the center of the trust region. In this case, this should be the optimum encountered so far.
    :param f: the function handle to evaluate x on (the acquisition function, in this case)
    :param max_hamming_dist: maximum Hamming distance from x_center
    :param step: number of maximum local search steps the algorithm is allowed to take.
    :return:
    """

    def _ls(hamming):
        """One restart of local search"""
        x_center_local = deepcopy(x_center)
        tol = step // 2
        trajectory = np.array([x_center_local])
        x = x_center_local

        acq_x = f(x).detach().cpu().numpy()
        for i in range(step):
            tol_ = tol
            is_valid = False
            while not is_valid:
                neighbour = sample_neighbour_ordinal(x, alphabet)
                if 0 < compute_hamming_dist_ordinal(x_center_local, neighbour) <= hamming \
                        and not any(np.equal(trajectory, neighbour).all(1)):
                    is_valid = True
                else:
                    tol_ -= 1
                if tol_ < 0:
                    print("Tolerance exhausted on this local search thread.")
                    return x, acq_x

            acq_x = f(x).detach().cpu().numpy()
            acq_neighbour = f(neighbour).detach().cpu().numpy()
            # print(acq_x, acq_neighbour)

            if acq_neighbour > acq_x:
                x = deepcopy(neighbour)
                trajectory = np.vstack((trajectory, deepcopy(x)))
        #print('local search thread ended with highest acquisition %s' % acq_x)
        return x, acq_x

    X = []
    fX = []
    for _ in tqdm(range(n_restart), desc='Local Search'):
        res = _ls(max_hamming_dist)
        X.append(res[0])
        fX.append(res[1])
        
    # select the top-1 biggest
    top_idices = np.argpartition(np.array(fX).flatten(), -1)[-1]  
    ret = np.array(X[top_idices]), np.array(fX).flatten()[top_idices]
    return ret
